import os
import threading
import requests
import m3u8
from Crypto.Cipher import AES
from concurrent.futures import ThreadPoolExecutor
import subprocess
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import shutil

class DownloaderApp:
    def __init__(self, root):
        self.root = root
        self.root.title("M3U8 视频下载器")
        self.root.geometry("600x400")

        # 变量存储
        self.save_path = tk.StringVar()
        self.save_name = tk.StringVar()
        self.m3u8_url = tk.StringVar()
        self.total_ts = 0
        self.completed_ts = 0
        self.progress = tk.DoubleVar()

        # 固定 TS 存储路径（项目目录下的 ts_segments）
        self.TS_DIR = os.path.join(os.path.dirname(__file__), "ts_segments")
        os.makedirs(self.TS_DIR, exist_ok=True)

        # GUI 组件
        self.create_widgets()

    def create_widgets(self):
        # M3U8 链接输入
        tk.Label(self.root, text="M3U8 链接:").pack(pady=5)
        self.m3u8_entry = tk.Entry(self.root, width=70, textvariable=self.m3u8_url)
        self.m3u8_entry.pack(pady=5)

        # 保存路径选择
        tk.Label(self.root, text="保存路径:").pack(pady=5)
        self.path_label = tk.Entry(self.root, width=60, textvariable=self.save_path)
        self.path_label.pack(pady=5)
        tk.Button(self.root, text="选择路径", command=self.select_save_path).pack(pady=5)

        # 保存文件名输入
        tk.Label(self.root, text="视频文件名（含后缀）:").pack(pady=5)
        self.name_entry = tk.Entry(self.root, width=60, textvariable=self.save_name)
        self.name_entry.pack(pady=5)

        # 开始下载按钮
        tk.Button(self.root, text="开始下载", command=self.start_download).pack(pady=10)

        # 进度条
        self.progress_bar = ttk.Progressbar(self.root, variable=self.progress, maximum=100, length=400)
        self.progress_bar.pack(pady=10)

        # 状态标签
        self.status_label = tk.Label(self.root, text="就绪")
        self.status_label.pack(pady=5)

    def select_save_path(self):
        path = filedialog.askdirectory()
        if path:
            self.save_path.set(path)

    def update_progress(self, value):
        self.progress.set(value)
        self.status_label.config(text=f"下载进度: {value:.1f}%")

    def start_download(self):
        url = self.m3u8_url.get().strip()
        save_path = self.save_path.get().strip()
        save_name = self.save_name.get().strip()

        if not all([url, save_path, save_name]):
            messagebox.showerror("错误", "请输入完整信息")
            return

        if not os.path.exists(save_path):
            messagebox.showerror("错误", "指定的保存路径不存在")
            return

        self.status_label.config(text="准备下载...")
        self.progress.set(0)
        self.progress_bar["style"] = "red.Horizontal.TProgressbar"

        # 启动下载线程
        threading.Thread(target=self.download_video, args=(url, save_path, save_name), daemon=True).start()

    def download_video(self, m3u8_url, save_path, save_name):
        try:
            # 解析 M3U8
            playlist = m3u8.load(m3u8_url)
            key_info, ts_list = self.parse_m3u8(playlist, m3u8_url)
            self.total_ts = len(ts_list)
            self.completed_ts = 0

            # 多线程下载 TS
            with ThreadPoolExecutor(max_workers=10) as executor:
                futures = []
                for i, ts_url in enumerate(ts_list):
                    future = executor.submit(self.download_ts, ts_url, i, key_info)
                    futures.append(future)

                # 等待所有线程完成
                for future in futures:
                    future.result()

            # 合并 TS
            output_mp4 = os.path.join(save_path, save_name)
            if self.merge_ts_files(output_mp4):
                self.progress.set(100)
                self.progress_bar["style"] = "green.Horizontal.TProgressbar"
                self.status_label.config(text="下载完成")
                self.open_folder(save_path)
                self.cleanup_ts()
            else:
                self.status_label.config(text="合并失败")
        except Exception as e:
            self.status_label.config(text="下载失败")
            messagebox.showerror("错误", str(e))

    def parse_m3u8(self, playlist, m3u8_url):  # 接收 m3u8_url 参数
        key = None
        if playlist.keys and playlist.keys[0]:
            key_url = playlist.keys[0].uri
            base_url = m3u8_url.rsplit("/", 1)[0]  # 使用传入的 m3u8_url
            if not key_url.startswith("http"):
                key_url = base_url + "/" + key_url
            try:
                key_response = requests.get(key_url)
                key_response.raise_for_status()
                key = (key_url, key_response.content)
            except requests.RequestException as e:
                print(f"获取密钥失败: {e}")
        ts_list = [seg.absolute_uri for seg in playlist.segments]
        return key, ts_list

    def download_ts(self, ts_url, index, key):
        try:
            headers = {"User-Agent": "Mozilla/5.0"}
            response = requests.get(ts_url, headers=headers, stream=True)
            response.raise_for_status()
            data = response.content

            if key:
                cipher = AES.new(key[1], AES.MODE_CBC, iv=key[1])
                data = cipher.decrypt(data)

            filename = os.path.join(self.TS_DIR, f"segment_{index:04d}.ts")
            with open(filename, "wb") as f:
                f.write(data)

            self.completed_ts += 1
            progress = (self.completed_ts / self.total_ts) * 100
            self.root.after(0, self.update_progress, progress)
        except Exception as e:
            print(f"下载错误 {ts_url}: {e}")

    def merge_ts_files(self, output_mp4):
        ts_files = sorted([os.path.join(self.TS_DIR, f) for f in os.listdir(self.TS_DIR) if f.endswith(".ts")])
        if not ts_files:
            print("没有找到 TS 文件")
            return False

        filelist_path = os.path.join(self.TS_DIR, "filelist.txt")

        with open(filelist_path, "w") as f:
            for ts in ts_files:
                f.write(f"file '{os.path.basename(ts)}'\n")

        command = [
            "ffmpeg",
            "-f", "concat",
            "-safe", "0",
            "-i", filelist_path,
            "-c", "copy",
            output_mp4
        ]
        try:
            # 设置工作目录为 TS 目录，确保 ffmpeg 找到文件
            subprocess.run(command, check=True, cwd=self.TS_DIR)
            return True
        except subprocess.CalledProcessError as e:
            print(f"FFmpeg 合并失败: {e}")
            return False

    def cleanup_ts(self):
        try:
            shutil.rmtree(self.TS_DIR)
            os.makedirs(self.TS_DIR)  # 重置空文件夹
        except Exception as e:
            print(f"删除 TS 目录失败: {e}")

    def open_folder(self, path):
        if os.name == "nt":
            os.startfile(path)
        elif os.name == "posix":
            subprocess.Popen(["xdg-open", path])

if __name__ == "__main__":
    root = tk.Tk()
    app = DownloaderApp(root)
    root.mainloop()